import torch
import torch.nn as nn
from .common import make_c2w, convert3x4_4x4

class LearnPose(nn.Module):
    def __init__(self, num_cams, learn_R, learn_t, init_c2w=None, num_rays=4096):
        """
        :param num_cams:
        :param learn_R:  True/False
        :param learn_t:  True/False
        :param cfg: config argument options
        :param init_c2w: (N, 4, 4) torch tensor
        """
        super(LearnPose, self).__init__()
        self.num_cams = num_cams
        self.init_c2w = None
        if init_c2w is not None:
            self.init_c2w = nn.Parameter(init_c2w, requires_grad=False)
        self.r = nn.Parameter(torch.zeros(size=(num_cams, 3), dtype=torch.float32), requires_grad=learn_R)  # (N, 3)
        self.t = nn.Parameter(torch.zeros(size=(num_cams, 3), dtype=torch.float32), requires_grad=learn_t)  # (N, 3)
        self.num_rays = num_rays
        # self.c2w = nn.Parameter(torch.eye(4, dtype=torch.float32).repeat(num_cams,1,1).reshape(-1, 16), requires_grad=True)

    def forward(self, cam_ids):
        # per image的训练方式
        # cam_index = cam_ids.unique().item()
        # r = self.r[cam_index]  # (3, ) axis-angle
        # t = self.t[cam_index]  # (3, )
        # c2w = make_c2w(r, t)  # (4, 4)
        # # learn a delta pose between init pose and target pose, if a init pose is provided
        # if self.init_c2w is not None:
        #     c2w = c2w @ self.init_c2w[cam_index]
        # return c2w.repeat(self.num_rays, 1, 1)
        #---------------------------------------
        # c2w_all = []
        # for iter, cam_index in enumerate(cam_ids):
        #     r = self.r[cam_index]  # (3, ) axis-angle
        #     t = self.t[cam_index]  # (3, )
        #     c2w_tmp = make_c2w(r, t)  # (4, 4)
        #     # learn a delta pose between init pose and target pose, if a init pose is provided
        #     if self.init_c2w is not None:
        #         c2w_tmp = c2w_tmp @ self.init_c2w[cam_index]
        #     c2w_all.append(c2w_tmp)
        # c2w = torch.cat(c2w_all)
        # return self.c2w.reshape(-1,4,4)[cam_ids]
        #----------------------------------------
        # 对上面程序进行优化加速  利用率还是有点低（但是比上面要好不少）
        c2w_all=[]
        for index, (r_value, t_value) in enumerate(zip(self.r, self.t)):
            c2w = make_c2w(r_value, t_value)
            if self.init_c2w is not None:
                c2w = c2w @ self.init_c2w[index]
            c2w_all.append(c2w.unsqueeze(0))
        c2w_tensor = torch.cat(c2w_all, dim=0)
        return c2w_tensor[cam_ids]
    # def forward(self, cam_ids):
    #     return self.c2w[cam_ids]

